import pandas as pd
import numpy as np
# from jqfactor import *
# from jqdata import *
import jqdatasdk
from jqdatasdk import *

jqdatasdk.auth("17320394612", "Fw023333")


# 获取数据主函数
# 输入股票池、指标名称、开始日期、结束日期
# 返回行标签为日期，列表签为股票名称的dataframe表格
def get_factor_data(stockPool, factor, date_start, date_end):
    # 获取股票池函数
    def get_stock(stockPool, begin_date):
        if stockPool == 'HS300':  # 用于获取沪深300股票池
            stockList = get_index_stocks('000300.XSHG', begin_date)
        elif stockPool == 'ZZ500':  # 用于获取中证500股票池
            stockList = get_index_stocks('399905.XSHE', begin_date)
        elif stockPool == 'ZZ800':  # 用于获取中证800股票池
            stockList = get_index_stocks('399906.XSHE', begin_date)
        elif stockPool == 'A':  # 用于获取全部A股股票池
            stockList = get_index_stocks('000002.XSHG', begin_date) + get_index_stocks('399107.XSHE', begin_date)
        else:  # 自定义输入股票池
            stockList = stockPool
        return stockList

        # 从财务库获取数据

    def get_factor_data1(factor, stock, date):
        if factor in val:
            q = query(valuation).filter(valuation.code.in_(stock))
            df = get_fundamentals(q, date)
        elif factor in bal:
            q = query(balance).filter(balance.code.in_(stock))
            df = get_fundamentals(q, date)
        elif factor in cf:
            q = query(cash_flow).filter(cash_flow.code.in_(stock))
            df = get_fundamentals(q, date)
        elif factor in inc:
            q = query(income).filter(income.code.in_(stock))
            df = get_fundamentals(q, date)
        elif factor in ind:
            q = query(indicator).filter(indicator.code.in_(stock))
            df = get_fundamentals(q, date)

        df.index = df['code']
        data = pd.DataFrame(index=df.index)
        data[date] = df[factor]

        return data.T

    # 获取日期列表
    date_list = get_trade_days(start_date=date_start, end_date=date_end)
    # 空df预备存储数据
    data = pd.DataFrame(columns=get_stock(stockPool, begin_date=date_list[-1]))

    # 获取五张财务基础所有指标名称
    val = get_fundamentals(query(valuation).limit(1)).columns.tolist()
    bal = get_fundamentals(query(balance).limit(1)).columns.tolist()
    cf = get_fundamentals(query(cash_flow).limit(1)).columns.tolist()
    inc = get_fundamentals(query(income).limit(1)).columns.tolist()
    ind = get_fundamentals(query(indicator).limit(1)).columns.tolist()
    all_columns = val + bal + cf + inc + ind

    # 循环时间列表获取指标数据
    for date in date_list:
        # 获取股票池
        all_stocks = get_stock(stockPool, date)
        # 获取因子数据
        if factor in all_columns:  # 可以从财务库直接取到因子值的因子
            data_temp = get_factor_data1(factor, all_stocks, date)
        else:  # 可以从因子库直接取到因子值的因子
            try:
                data_temp = get_factor_values(all_stocks, [factor], end_date=date, count=1)[factor]
            except:
                print('系统暂不能获取该因子，请获取其他因子')
                break
        data = pd.concat([data, data_temp], axis=0)
    return data


# 示例，获取沪深300成分股在2019年前两个月的市盈率数据
res = get_factor_data('HS300', 'pe_ratio', '2024-06-10', '2024-09-15')
print(res)
